"""
outputs = model(input_ids, attention_mask=attention_mask, labels=labels) 
# labels是真实标签和input_ids完全相等, 这个代码与之自动计算交叉熵损失
# 但是这个代码input_ids和labels形状完全相同啊是如何计算的交叉熵损失呢？
因为inputs_ids会被底层的nn.embedding等一系列网络层升维,变成维度Crossentropy([length,modeldim],[length])
loss = outputs.loss
"""

# 库
from transformers import GPT2Tokenizer, GPT2LMHeadModel
from transformers import Trainer, TrainingArguments, AdamW
from torch.utils.data import DataLoader
import torch
import json
from typing import Dict
from torch.nn import CrossEntropyLoss
# openai gpt2 可以在huggingface上看看
# 加载预训练好的基本模型，同时把模型放置到gpu上面
tokenizer = GPT2Tokenizer.from_pretrained('gpt2',local_files_only=True)  # downloaded in default
model = GPT2LMHeadModel.from_pretrained('gpt2',local_files_only=True)
device="cuda:0"
#model.to(device)
# 微调数据集的构建
with open("./data.jsonl", 'r', encoding='utf-8') as f:
     data=f.readlines()
#大多数llm都需要设置padding的token如果没有设置会报错，由于vocab.json：https://hf-mirror.com/openai-community/gpt2/tree/main 中是有eos_tokens的所以使用这个     
print(tokenizer.eos_token)
tokenizer.pad_token = tokenizer.eos_token

class Dataset(torch.utils.data.Dataset):
    def __init__(self, data:list[Dict]):
        self.data = data
   
    def __getitem__(self, idx):
        mid=json.loads(self.data[idx])
        text = mid["Q"] + tokenizer.eos_token + mid["A"]#中间的这个字符是因为在openai提供的仓库中vocab.json本来就有的哈
        inputs = tokenizer.encode_plus(text, 
        max_length=32,  # 设置最大序列长度
        padding='max_length',  # 进行填充以确保序列达到最大长度
        truncation=True,  # 截断超出最大长度的序列
        return_attention_mask=True,  # 返回注意力掩码
        return_tensors='pt')
        attention_mask = inputs['attention_mask']
        #inputs.to(device)
        return {
            'input_ids': inputs['input_ids'],
            'attention_mask':attention_mask,   #0表示发生填充的地方,1表示没有填充的地方
            'labels': inputs['input_ids']
        }
    
        """
        theory explain: different to pre_train, SFT is not predicting next words, instead, it's to give the same sequence
        Q<sep>A,
        """
        # trainer类的数据格式应该如何构建
    def __len__(self):
        return len(self.data)

dataset=Dataset(data)
print(model(**dataset[0]))   #检查模型是否可以正常地推理
dataloader = DataLoader(dataset, batch_size=2, shuffle=True)

# 微调设置
# 定义损失函数和优化器
optimizer = AdamW(model.parameters(), lr=5e-5)

# 训练模型
model.train()
for epoch in range(3):  # 训练3个epoch
    for step, batch in enumerate(dataloader):
        input_ids = batch['input_ids']
        attention_mask = batch['attention_mask']
        labels = batch['labels']
        print(input_ids)
        # 前向传播
        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
        assert False
        loss = outputs.loss

        # 反向传播和优化
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        print(f'Epoch: {epoch}, Step: {step}, Loss: {loss.item()}')

# 7. 保存模型
model.save_pretrained('./SFTgpt')
tokenizer.save_pretrained('./SFTgpt')